"""
Author: Morphlng
Date: 2024-01-23 15:26:43
LastEditTime: 2025-02-25 15:46:13
LastEditors: Morphlng
Description: Helper functions for carla adapter
FilePath: /DrivingGym/src/driving_gym/simulation/adapter/carla/utils/misc.py
"""

from __future__ import annotations

import logging
import math
import os
import random
import signal
import socket
import subprocess
import time

import carla

from driving_gym.simulation.common.geom import (
    BoundingBox,
    Location,
    Rotation,
    Transform,
    Vector3D,
)

logger = logging.getLogger(__name__)


def start_local_server(
    port: int, resolution: tuple[int, int] = (800, 600), offscreen: bool = False
):
    executable = os.environ.get("CARLA_SERVER", None)
    if executable is None or not os.path.exists(executable):
        raise FileNotFoundError(
            "Make sure CARLA_SERVER environment"
            " variable is set & is pointing to the"
            " CARLA server startup script (Carla"
            "UE4.sh/exe). Refer to the README file/docs."
        )

    is_windows_platform = os.name == "nt"
    process = None
    if offscreen:
        process = subprocess.Popen(
            f"{executable} -RenderOffScreen -benchmark -fps=20 -carla-server -world-port={port} -carla-streaming-port=0",
            shell=True,
            # for Linux
            preexec_fn=None if is_windows_platform else os.setsid,
            # for Windows (not necessary)
            creationflags=(
                subprocess.CREATE_NEW_PROCESS_GROUP if is_windows_platform else 0
            ),
            stdout=subprocess.DEVNULL,
            stderr=subprocess.STDOUT,
        )
        logger.info("Running CARLA server in headless mode")
    else:
        process = subprocess.Popen(
            [
                executable,
                "-windowed",
                "-ResX=",
                str(resolution[0]),
                "-ResY=",
                str(resolution[1]),
                "-benchmark",
                "-fps=20",
                "-carla-server",
                f"-carla-rpc-port={port}",
                "-carla-streaming-port=0",
            ],
            # for Linux
            preexec_fn=None if is_windows_platform else os.setsid,
            # for Windows (not necessary)
            creationflags=(
                subprocess.CREATE_NEW_PROCESS_GROUP if is_windows_platform else 0
            ),
            stdout=subprocess.DEVNULL,
            stderr=subprocess.STDOUT,
        )
        logger.info("Running simulation in single-GPU mode")

    time.sleep(3)
    return process


def stop_local_server(process: subprocess.Popen):
    if os.name == "nt":
        # for Windows
        subprocess.call(["taskkill", "/F", "/T", "/PID", str(process.pid)])
    else:
        # for Linux
        pgid = os.getpgid(process.pid)
        os.killpg(pgid, signal.SIGKILL)


def get_tcp_port(port: int = 0) -> int:
    """Get a free tcp port number

    Args:
        port: (default 0) port number. When `0` it will be assigned a free port dynamically

    Return:
        a port number requested if free otherwise an unhandled exception would be thrown
    """

    s = socket.socket()
    s.bind(("", port))
    server_port = s.getsockname()[1]
    s.close()
    return server_port


def localize_transform(transform: carla.Transform):
    """Covert carla.Transform into Transform"""

    return Transform(
        location=Location(
            x=transform.location.x,
            y=transform.location.y,
            z=transform.location.z,
        ),
        rotation=Rotation(
            pitch=transform.rotation.pitch,
            yaw=transform.rotation.yaw,
            roll=transform.rotation.roll,
        ),
    )


def localize_location(location: carla.Location):
    """Covert carla.Location into Location"""

    return Location(
        x=location.x,
        y=location.y,
        z=location.z,
    )


def localize_rotation(rotation: carla.Rotation):
    """Covert carla.Rotation into Rotation"""

    return Rotation(
        pitch=rotation.pitch,
        yaw=rotation.yaw,
        roll=rotation.roll,
    )


def localize_vector(vector: carla.Vector3D):
    """Covert carla.Vector3D into Vector3D"""

    return Vector3D(
        x=vector.x,
        y=vector.y,
        z=vector.z,
    )


def localize_bounding_box(box: carla.BoundingBox):
    """Covert carla.BoundingBox into BoundingBox"""

    return BoundingBox(
        extent=Vector3D(
            x=box.extent.x,
            y=box.extent.y,
            z=box.extent.z,
        ),
        location=Location(
            x=box.location.x,
            y=box.location.y,
            z=box.location.z,
        ),
        rotation=Rotation(
            pitch=box.rotation.pitch,
            yaw=box.rotation.yaw,
            roll=box.rotation.roll,
        ),
    )


def carla_tranform(transform: Transform):
    """Convert Transform into carla.Transform"""

    return carla.Transform(
        location=carla.Location(
            x=transform.location.x,
            y=transform.location.y,
            z=transform.location.z,
        ),
        rotation=carla.Rotation(
            pitch=transform.rotation.pitch,
            yaw=transform.rotation.yaw,
            roll=transform.rotation.roll,
        ),
    )


def carla_location(location: Location):
    """Convert Location into carla.Location"""

    return carla.Location(
        x=location.x,
        y=location.y,
        z=location.z,
    )


def carla_rotation(rotation: Rotation):
    """Convert Rotation into carla.Rotation"""

    return carla.Rotation(
        pitch=rotation.pitch,
        yaw=rotation.yaw,
        roll=rotation.roll,
    )


def carla_vector(vector: Vector3D):
    """Convert Vector3D into carla.Vector3D"""

    return carla.Vector3D(
        x=vector.x,
        y=vector.y,
        z=vector.z,
    )


def vehicle_control(control: dict):
    """Convert control dict into carla.VehicleControl"""

    return carla.VehicleControl(
        throttle=control.get("throttle", 0.0),
        steer=control.get("steer", 0.0),
        brake=abs(control.get("brake", 0.0)),
        hand_brake=control.get("hand_brake", False),
        reverse=control.get("reverse", False),
        manual_gear_shift=control.get("manual_gear_shift", False),
        gear=control.get("gear", 0),
    )


def walker_control(control: dict):
    """Convert control dict into carla.WalkerControl"""

    return carla.WalkerControl(
        speed=control.get("speed", 0.0),
        direction=control.get("direction", carla.Vector3D(0, 0, 0)),
        jump=control.get("jump", False),
    )


def parse_random_range(
    random_range: "tuple[float, float] | list | float",
) -> tuple[float, float]:
    """Parse the random range from the given input.

    Args:
        random_range: A tuple or list with two elements defining the lower and upper bound,
        or a float defining the symmetric range around zero.

    Returns:
        lower_bound, upper_bound (tuple).
    """
    if isinstance(random_range, (tuple, list)):
        return tuple(random_range)
    else:
        return (-random_range, random_range)


def get_location_from_dict(loc_conf: dict, carla_map: carla.Map) -> tuple[float, ...]:
    """Parse a location configuration dictionary and return the corresponding coordinates.

    Args:
        loc_conf (dict): location configuration dictionary
        carla_map (carla.Map): Carla map

    Raises:
        ValueError: Unsupported location configuration

    Returns:
        location (tuple[float, ...]): The corresponding coordinates
    """
    h = loc_conf.get("h", None)
    yaw = loc_conf.get("yaw", math.degrees(h) if h is not None else None)
    z = loc_conf.get("z", 2.0)

    # Handle OpenDrive location specification
    if {"road_id", "lane_id", "s"}.issubset(loc_conf):
        road_id, lane_id, s = (
            loc_conf["road_id"],
            loc_conf["lane_id"],
            loc_conf["s"],
        )

        # Apply random offset to 's' if specified
        if "random_range_s" in loc_conf:
            range_value_s = loc_conf["random_range_s"]
            s = random.uniform(range_value_s[0], range_value_s[1])
        elif "random_s" in loc_conf:
            lower, upper = parse_random_range(loc_conf["random_s"])
            s = max(0.1, s + random.uniform(lower, upper))

        waypoint = carla_map.get_waypoint_xodr(road_id, lane_id, int(s))
        x, y, z = (
            waypoint.transform.location.x,
            waypoint.transform.location.y,
            max(z, waypoint.transform.location.z),
        )

        dh = loc_conf.get("dh", None)
        dyaw = loc_conf.get("dyaw", math.degrees(dh) if dh is not None else None)
        if dyaw is not None:
            yaw = waypoint.transform.rotation.yaw + dyaw
    # Handle absolute position
    elif "x" in loc_conf and "y" in loc_conf:
        x, y = loc_conf["x"], loc_conf["y"]
    # If a nested 'location' is provided, recursively parse the location
    elif "location" in loc_conf:
        location = parse_location(loc_conf["location"], carla_map)
        x, y, z = location[0], location[1], max(location[2], z)
        if len(location) > 3:
            yaw = location[3]
    else:
        raise ValueError(
            "Dictionary location configuration must contain one of: 'location' /"
            "'road_id', 'lane_id', 's' / 'x' and 'y' keys."
        )

    # Apply random offsets if specified
    variables = {"x": x, "y": y, "z": z, "yaw": yaw}
    for axis in variables.keys():
        random_key = f"random_{axis}"
        random_range_key = f"random_range_{axis}"

        if random_range_key in loc_conf:
            range_value = loc_conf[random_range_key]
            variables[axis] = random.uniform(range_value[0], range_value[1])
        elif random_key in loc_conf:
            lower, upper = parse_random_range(loc_conf[random_key])
            variables[axis] += random.uniform(lower, upper)

    x, y, z, yaw = variables.values()
    return (x, y, z) if yaw is None else (x, y, z, yaw)


def parse_location(
    loc_conf: "int | tuple[float, ...] | list | dict", carla_map: carla.Map
) -> Transform:
    """Parse a location configuration and return the corresponding coordinates.

    Args:
        loc_conf (int | tuple[float, ...] | list | dict): location configuration
        carla_map (carla.Map): map in which the location is defined

    Raises:
        ValueError: Unsupported location configuration

    Returns:
        location (Transform): The corresponding coordinates
    """
    if isinstance(loc_conf, (carla.Transform, Transform)):
        return loc_conf
    elif isinstance(loc_conf, (tuple, list)):
        loc = tuple(loc_conf)
    elif isinstance(loc_conf, dict):
        loc = get_location_from_dict(loc_conf, carla_map)
    else:
        raise ValueError(f"Unsupported location configuration: {loc_conf}")

    if len(loc) == 3:
        loc = (*loc, 0.0)
    return Transform(Location(*loc[:3]), Rotation(yaw=loc[3]))
